package exp1dtree;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;

import exp1dtree.dtree.DTree;
import exp1dtree.record.Record;

public final class Driver {

	private static Driver driver;
	
	private Option option;
	private DTree dtree;
	
	private List<Record> trainList;
	private List<Record> validationList;
	private List<Record> testList;

	public static void main(String[] args) throws IOException {
		driver = new Driver();
		driver.option = new Option(args);
		driver.init();
		driver.train();
		driver.test();
	}
	
	public void init() throws IOException {
		Record.init();
		
		// Read train set
		BufferedReader br = new BufferedReader(
				new InputStreamReader(driver.option.getTrainInput()));
		String s;
		testList = new ArrayList<Record>();
		while ((s = br.readLine()) != null) {
			Record aRecord = new Record(s);
			testList.add(aRecord);
		}
		
		// random shuffle
		java.util.Collections.shuffle(testList, new Random());
		int trainTot = testList.size() * option.getFraction() / 100;
		int trainPos = trainTot / (1 + option.getRatio());
		int trainNeg = trainTot - trainPos;
		int validation = testList.size() * option.getValidation() / 100;
		if (testList.size() - trainTot < validation)
			validation = testList.size() - trainTot;
		trainList = new ArrayList<Record>();
		validationList = new ArrayList<Record>();
		for (Record r : testList) {
			if (trainTot > 0) {
				if (r.isPositive()) {
					if (trainPos > 0) {
						--trainTot;
						--trainPos;
						trainList.add(r);
						continue;
					}
				} else if (trainNeg > 0) {
					--trainTot;
					--trainNeg;
					trainList.add(r);
					continue;
				}
			}
			if (validation > 0) {
				--validation;
				validationList.add(r);
			}
		}
		
		// Read test set
		br = new BufferedReader(
				new InputStreamReader(option.getTestInput()));
		testList.clear();
		while ((s = br.readLine()) != null) {
			Record aRecord = new Record(s);
			testList.add(aRecord);
		}
	}

	public void train() {
		DTree.trainList = trainList;
		DTree.validationList = validationList;
		DTree.able = new boolean[Record.attCnt];
		for (int i = 0; i < DTree.able.length; ++i)
			DTree.able[i] = true;
		dtree = new DTree();
		dtree.build(0, trainList.size(), option.getStrategy());
		if (option.getStrategy() == DTree.Strategy.POST) {
			option.getOutput().println("Before pruning:");
			test();
			dtree.postPrune(0, validationList.size());
			option.getOutput().println("After pruning:");
		}
	}
	
	public void runTest(List<Record> list) {
		TestResult rst = new TestResult();
		for (Record r : list)
			rst.add(dtree.simulate(r), r.isPositive());
		option.getOutput().print(rst);
	}
	
	public void test() {
		option.getOutput().println("Tree size:\t" + dtree.getSize());
		option.getOutput().println("Training Result:");
		runTest(trainList);
		option.getOutput().println("Testing Result:");
		runTest(testList);
	}

}
